#!/usr/bin/env python
# encoding: utf-8

from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy import INTEGER, VARCHAR, ARRAY, TIMESTAMP, Column, create_engine, ForeignKey, TEXT
from sqlalchemy import JSON, DATE, PrimaryKeyConstraint, REAL, BIGINT, FLOAT, UniqueConstraint
from sqlalchemy.orm import sessionmaker, relationship

ORMBase = declarative_base()
dbengine = create_engine(
    'postgresql+psycopg2://postgres:postgres@10.0.0.39:5432/ass',
    echo=True, client_encoding='utf8'
)
Session = sessionmaker(bind=dbengine)


def SQL(arg):
    def _deco(func):
        def __deco(*args, **kwargs):
            _db = Session()
            kwargs['db'] = _db
            ret = func(*args, **kwargs)
            if "Insert" in arg:
                assert isinstance(ret, list), f"{func} must return list."
                _db.add_all(ret)
            _db.commit()
            _db.close()

            return ret
        return __deco
    return _deco


class Quota(ORMBase):
    __tablename__ = 't_quota'
    id = Column(INTEGER, primary_key=True, autoincrement=True)
    name = Column(VARCHAR)
    formula = Column(VARCHAR)
    create_time = Column(TIMESTAMP)
    status = Column(INTEGER)
    sources = Column(ARRAY(VARCHAR))
    market = Column(INTEGER)

    def __repr__(self):
        return f"< Quota ({self.create_time}){self.id}:{self.name}={self.formula} >"


class Report(ORMBase):
    __tablename__ = 't_report'
    __table_args__ = (
        PrimaryKeyConstraint('code', 'date'),
        UniqueConstraint('code', 'date'),
    )

    code = Column(VARCHAR)
    date = Column(DATE)
    financial = Column(JSON)

    def __repr__(self):
        return f"< Report ({self.date}){self.code}:{self.financial.keys()} >"


class Strategy(ORMBase):
    __tablename__ = 't_strategy'

    id = Column(INTEGER, primary_key=True, autoincrement=True, nullable=False)
    name = Column(VARCHAR, nullable=False)
    formulas = Column(ARRAY(VARCHAR), nullable=False)

    def __repr__(self):
        formula = ""
        index = 1
        for i in self.formulas:
            formula += f"\n\t{index}\t{i}"
        return f"< Strategy {self.id}{self.name} {formula} >"


class StockDay(ORMBase):
    __tablename__ = 't_stock_day'
    __table_args__ = (
        PrimaryKeyConstraint('code', 'date'),
        UniqueConstraint('code', 'date'),
    )

    code = Column(VARCHAR)
    date = Column(DATE)
    close = Column(REAL)
    high = Column(REAL)
    low = Column(REAL)
    open = Column(REAL)
    chg = Column(REAL)
    pchg = Column(REAL)
    turnover = Column(REAL)
    voturnover = Column(BIGINT)
    vaturnover = Column(FLOAT)
    tcap = Column(FLOAT)
    mcap = Column(FLOAT)

    def __repr__(self):
        return f"""
< StockDay {self.code} {self.date}: 
    开盘价: {self.open}, \t收盘价: {self.close}      
    最高价: {self.high}, \t最低价: {self.low}
    涨跌额: {self.chg},  \t涨跌幅: {self.pchg}        
    换手率: {self.turnover}, \t成交量: {self.voturnover}, \t成交金额: {self.vaturnover}
    总市值:   {self.tcap}
    流程市值: {self.mcap} >"""

    def get_market(self):
        return ",".join(map(lambda x: str(x), [
            self.code, self.open, self.close, self.close,
            self.high, self.low, self.close, self.close,
            self.voturnover, self.vaturnover, self.date.strftime("%Y-%m-%d")]))


class StockInfo(ORMBase):
    __tablename__ = 't_stockinfo'

    code = Column(VARCHAR, primary_key=True)
    name = Column(VARCHAR)
    market = Column(VARCHAR)
    profession = Column(ARRAY(VARCHAR))

    def __repr__(self):
        return f"< StockInfo {self.code} {self.name} {'、'.join(self.profession)} >"


class ReportQuota(ORMBase):
    __tablename__ = 't_report_quota'
    id = Column(INTEGER, primary_key=True, autoincrement=True)
    name = Column(VARCHAR, unique=True)
    codes = Column(ARRAY(VARCHAR))
    coverage = Column(FLOAT)

    def __repr__(self):
        return f"< ReportQuota {self.id}:{self.name}({self.coverage})\t{'、'.join(self.codes)}) >"


class Industry(ORMBase):
    __tablename__ = "t_industry"

    id = Column(INTEGER, primary_key=True, autoincrement=True)
    name = Column(VARCHAR, unique=True)
    codes = Column(ARRAY(VARCHAR))

    def __repr__(self):
        return f"< Industry {self.id}:{self.name}\t{'、'.join(self.codes)}) >"


class NewIndustry(ORMBase):
    __tablename__ = "t_new_industry"

    id = Column(INTEGER, primary_key=True, autoincrement=True)
    name = Column(VARCHAR, unique=True)
    codes = Column(ARRAY(VARCHAR))

    def __repr__(self):
        return f"< New Industry {self.id}:{self.name}\t{'、'.join(self.codes)}) >"


class MarketKeyWord(ORMBase):
    __tablename__ = "t_market_keyword"

    id = Column(INTEGER, primary_key=True, autoincrement=True)
    name = Column(VARCHAR, unique=True, nullable=True)
    pos = Column(INTEGER, nullable=True)

    def __repr__(self):
        return f"< MarketKeyWord {self.id}:{self.name}:{self.pos}) >"


class StrategyRuntime(ORMBase):
    __tablename__ = 't_strategy_runtime'

    id = Column(INTEGER, primary_key=True, autoincrement=True)
    strategy_id = Column(INTEGER, ForeignKey('t_strategy.id'))
    code = Column(VARCHAR)
    status = Column(INTEGER)
    schedule = Column(FLOAT)
    create_time = Column(TIMESTAMP)
    finished_time = Column(TIMESTAMP)
    change_by = Column(VARCHAR)
    success = Column(ARRAY(TEXT))
    failed = Column(ARRAY(TEXT))
    quota_result = Column(JSON, default={})

    strategy = relationship("Strategy", backref="t_strategy_runtime")

    def __repr__(self):
        return f"""< StrategyRuntime 
    {self.id}:{self.code}:{self.schedule} 
    Signal:{self.change_by} 
    Success: {",".join(self.success)}
    Failed: {",".join(self.failed)}
>"""


@SQL("Select")
def unit_class(db):
    for i in db.query(Quota).limit(5):
        print(i)

    for i in db.query(Report).limit(5):
        print(i)

    for i in db.query(StockDay).limit(5):
        print(i)

    for i in db.query(StockInfo).limit(5):
        print(i)

    for i in db.query(ReportQuota).limit(5):
        print(i)

    for i in db.query(Industry).limit(5):
        print(i)

    for i in db.query(StrategyRuntime).limit(5):
        print(i, type(i.success), isinstance(i.success, list))

    print('-'*100)
    r = Session()
    a = r.execute("select id, '@'||name from t_quota;")
    for i in a:
        print(i)
    print('#'*10, r.query(StockDay).limit(1).all()[0].get_market())
    print(r.query(Strategy).all()[0].formulas)
    import random
    print(random.sample(r.query(StockInfo.code).order_by(StockInfo.code).all(), 10))
    print(r.query(StockDay).filter(StockDay.date >= '2018-06-05').limit(10).all())
    from sqlalchemy.sql import func
    import datetime
    d = r.query(func.min(StockDay.date)).filter(StockDay.code == '000797').one()[0]+datetime.timedelta(days=55)
    print(d, type(d))
    r.close()


def create_all():
    ORMBase.metadata.create_all(dbengine)


if __name__ == "__main__":
    create_all()
    unit_class()
    # from source.Lib import Tools
    # print(dir(Tools))
    # print(Tools.__dict__['unit'])
